package com.artisan.day01;

import com.artisan.common.ApiKeys;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.mapdb.DB;
import org.mapdb.DBMaker;

import java.util.List;
import java.util.Map;

import static dev.langchain4j.data.message.ChatMessageDeserializer.messagesFromJson;
import static dev.langchain4j.data.message.ChatMessageSerializer.messagesToJson;
import static dev.langchain4j.model.openai.OpenAiChatModelName.GPT_4_O_MINI;
import static org.mapdb.Serializer.INTEGER;
import static org.mapdb.Serializer.STRING;

public class ServiceWithPersistentMemoryForEachUserExample {

    interface Assistant {

        String chat(@MemoryId int memoryId, @UserMessage String userMessage);
    }

    /**
     * 如何使用聊天AI助手进行对话
     * 它首先设置了一个持久化的聊天记忆存储，然后配置了一个聊天记忆提供者
     * 接着，它创建了一个聊天语言模型，并将其与记忆提供者一起用于构建一个AI助手
     * 最后，它展示了如何与AI助手进行聊天，并记住对话上下文
     */
    public static void main(String[] args) {

        // 创建一个持久化的聊天记忆存储实例
        PersistentChatMemoryStore store = new PersistentChatMemoryStore();

        // 定义一个聊天记忆提供者，它根据记忆ID构建消息窗口聊天记忆
        // 这里设定了每个对话最多保留10条消息
        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(10)
                .chatMemoryStore(store)
                .build();

        // 创建一个聊天语言模型实例，需要指定API密钥和模型名称
        ChatLanguageModel model = OpenAiChatModel.builder()
                .apiKey(ApiKeys.OPENAI_API_KEY)
                .modelName(GPT_4_O_MINI)
                .build();

        // 使用上述模型和聊天记忆提供者构建一个AI助手实例
        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(model)
                .chatMemoryProvider(chatMemoryProvider)
                .build();

        // 与AI助手进行聊天，并打印助手的回复
        // 这里是第一次聊天，分别与ID为1和2的对话对象打招呼
//        System.out.println(assistant.chat(1, "你好，我的名字叫小工匠"));
//        System.out.println(assistant.chat(2, "你好，我的名字叫Artisan"));

        // 现在，注释掉上面的两行代码，取消下面两行代码的注释，然后再次运行。
        // 这样做是为了演示AI助手如何根据之前的对话记住用户的名字
         System.out.println(assistant.chat(1, "我是谁?"));
         System.out.println(assistant.chat(2, "我是谁?"));
    }

    /**
     * 持久化聊天内存存储类
     * 实现了ChatMemoryStore接口，用于持久化存储聊天消息
     * 该类使用MapDB库创建一个持久化的数据库，用于存储聊天消息
     */
    static class PersistentChatMemoryStore implements ChatMemoryStore {

        // 创建一个持久化的数据库实例
        private final DB db = DBMaker.fileDB("multi-user-chat-memory.db").transactionEnable().make();
        // 创建一个整数到字符串的哈希映射，用于存储聊天消息
        private final Map<Integer, String> map = db.hashMap("messages", INTEGER, STRING).createOrOpen();

        /**
         * 根据memoryId获取聊天消息
         *
         * @param memoryId 聊天消息的标识符
         * @return 聊天消息的列表，如果memoryId不存在，则返回null
         */
        @Override
        public List<ChatMessage> getMessages(Object memoryId) {
            String json = map.get((int) memoryId);
            return messagesFromJson(json);
        }

        /**
         * 更新特定memoryId的聊天消息
         *
         * @param memoryId 聊天消息的标识符
         * @param messages 要更新的聊天消息列表
         */
        @Override
        public void updateMessages(Object memoryId, List<ChatMessage> messages) {
            String json = messagesToJson(messages);
            map.put((int) memoryId, json);
            db.commit();
        }

        /**
         * 删除特定memoryId的聊天消息
         *
         * @param memoryId 聊天消息的标识符
         */
        @Override
        public void deleteMessages(Object memoryId) {
            map.remove((int) memoryId);
            db.commit();
        }
    }
}